#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import json
from typing import List

from agentscope.agent import AgentBase
from agentscope.message import Msg
from agentscope.model import DashScopeChatModel
from agentscope.formatter import DashScopeChatFormatter
from agentscope.memory import InMemoryMemory
from pydantic import BaseModel, Field

from src.utils.env_config import aliyun_model_config


# Pydantic模型用于结构化输出
class KeywordExtractionResult(BaseModel):
    """关键词提取结果"""
    keywords: List[str] = Field(description="提取的关键词列表")

class KeywordExtractorAgent(AgentBase):
    """T1 - 关键词提取智能体"""

    def __init__(self):
        super().__init__()
        self.name = "KeywordExtractorAgent"
        self.sys_prompt = """你是一个专业的关键词提取智能体。你的任务是从给定的报告主题中精准提取出核心关键词或短语。
        要求：
        1. 分析输入主题的语言结构和语义重点；
        2. 提取多个最具代表性的关键词或短语；
        3. 关键词应为名词性短语，能独立表达一个核心概念；
        4. 输出格式必须为JSON，仅包含一个字段"keywords"，值为字符串数组；
        5. 不要解释过程，不要添加额外内容。
        示例输入："2025 Q3 AI 行业趋势"
        示例输出：{"keywords": ["2025 Q3", "AI 行业", "趋势"]}"""

        self.model = DashScopeChatModel(
            model_name=aliyun_model_config["model_name"],
            api_key=aliyun_model_config["api_key"],
            stream=False
        )
        self.formatter = DashScopeChatFormatter()
        self.memory = InMemoryMemory()

    async def reply(self, msg: Msg) -> Msg:
        """处理关键词提取请求"""
        await self.memory.add(msg)

        # 准备提示
        prompt = await self.formatter.format([
            Msg("system", self.sys_prompt, "system"),
            msg
        ])

        # 调用模型进行结构化输出
        response = await self.model(prompt)

        # 创建回复消息
        reply_msg = Msg(
            name=self.name,
            content=response.content,
            role="assistant"
        )

        await self.memory.add(reply_msg)
        return reply_msg

    async def extract_keywords(self, topic: str) -> List[str]:
        """提取关键词"""
        user_msg = Msg(
            name="user",
            content=f"如果输入主题中包含'技术突破'、'创新案例'、'政策法规'词语时，提取的关键词列表一定要包含这些词语。现在开始处理：\n输入主题：{topic}",
            role="user"
        )

        response_msg = await self.reply(user_msg)

        try:
            # 提取文本内容
            content_text = ""
            if isinstance(response_msg.content, list):
                for block in response_msg.content:
                    if hasattr(block, 'text'):
                        content_text += block.text
                    elif isinstance(block, dict) and 'text' in block:
                        content_text += block['text']
            else:
                content_text = str(response_msg.content)

            result = json.loads(content_text)
            return result.get("keywords", [])
        except json.JSONDecodeError:
            # 如果JSON解析失败，返回默认关键词
            print(f"JSON解析失败:{response_msg.content}")
            return [topic]

    async def observe(self, msg: Msg) -> None:
        """观察消息"""
        await self.memory.add(msg)

    async def handle_interrupt(self) -> Msg:
        """处理中断"""
        return Msg(
            name=self.name,
            content="关键词提取任务被中断",
            role="assistant"
        )
